package net.youdi.gateway.filter;

import com.fasterxml.jackson.core.json.JsonReadFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import net.youdi.gateway.context.BodyHandlerFunction;
import net.youdi.gateway.context.BodyHandlerServerHttpResponseDecorator;
import net.youdi.gateway.context.DefaultClientResponseAdapter;
import net.youdi.gateway.context.GatewayLog;
import net.youdi.gateway.handler.GateWayRequestCacheHandler;
import net.youdi.gateway.handler.request.RequestBodyHandler;
import net.youdi.gateway.handler.response.RewriteResponseHandlerFunction;
import net.youdi.gateway.options.FilterOrderEnum;
import net.youdi.gateway.properties.GatewayProperties;
import net.youdi.gateway.util.GateWayUtil;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.filter.factory.rewrite.CachedBodyOutputMessage;
import org.springframework.cloud.gateway.filter.headers.HttpHeadersFilter;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ReactiveHttpOutputMessage;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.util.MultiValueMap;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.reactive.function.server.ServerResponse;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.UnsupportedEncodingException;
import java.net.URI;
import java.net.URLEncoder;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.ORIGINAL_RESPONSE_CONTENT_TYPE_ATTR;

/**
 * Gateway Context Filter
 * @description 将请求数据缓存到上下文中
 */
@Slf4j
public class GatewayRequestFilter implements GlobalFilter, Ordered {
    public static final String GLOBAL_LOG_ATTR = "gateway_log";

    @Autowired
    private volatile List<HttpHeadersFilter> headersFilters;

    //重写响应数据处理器
    @Autowired
    private RewriteResponseHandlerFunction rewriteResponseHandler;
    //请求数据处理器
    @Autowired(required = false)
    private List<RequestBodyHandler> requestBodyHandlers;

    private final ObjectMapper objectMapper;

    private final GatewayProperties gatewayProperties;

    private GateWayRequestCacheHandler gateWayRequestCacheHandler;

    public GatewayRequestFilter(ObjectMapper objectMapper, GatewayProperties gatewayProperties,
                                GateWayRequestCacheHandler gateWayRequestCacheHandler) {
        this.objectMapper = objectMapper;
        this.objectMapper.configure(JsonReadFeature.ALLOW_UNESCAPED_CONTROL_CHARS.mappedFeature(),true);
        this.gatewayProperties = gatewayProperties;
        this.gateWayRequestCacheHandler = gateWayRequestCacheHandler;
    }

    @Override
    public int getOrder() {
        return FilterOrderEnum.REQUEST_LOG_FILTER.getOrder()+5;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
      //如果不记录日志，直接放行
        if(!gatewayProperties.getLogRequest()){
            log.debug("[GatewayContext]Properties Set To Not Log Request Data");
            return chain.filter(exchange);
        }
        //记录相关日志
        long start = System.currentTimeMillis();
        GatewayLog gatewayLog = new GatewayLog();
        exchange.getAttributes().put(GLOBAL_LOG_ATTR,gatewayLog);
        gatewayLog.setStartTime(start);
        gatewayLog.setUserInfo(exchange.getAttributes().get(GatewayProperties.USER_KEY));
        return handleRequest(exchange, chain, gatewayLog).then(Mono.fromRunnable(() -> {
            gatewayLog.setEndTime(System.currentTimeMillis());
            if (gatewayProperties.getLogCache() && gateWayRequestCacheHandler != null) {
                gateWayRequestCacheHandler.cache(gatewayLog);
            }
        }));

    }

    /**
     * log request
     * @param exchange
     * @param chain
     */
    private Mono<Void> handleRequest(ServerWebExchange exchange, GatewayFilterChain chain, GatewayLog gatewayLog){

        ServerHttpRequest request = exchange.getRequest();
        URI requestURI = request.getURI();
        HttpHeaders headers = request.getHeaders();
        String ipAddress = GateWayUtil.getIpAddress(request);
        MultiValueMap<String, String> queryParams = request.getQueryParams();
        gatewayLog.setRequestMethod(request.getMethod()).setRequestUrl(requestURI)
                .setRequestHeaders(headers).setRequestIp(ipAddress).setQueryParams(queryParams);
        MediaType contentType = headers.getContentType();
        long length = headers.getContentLength();
        if(length>0 && null != contentType && contentType.includes(MediaType.APPLICATION_JSON)){
            return handleRequestBody(exchange,chain,gatewayLog);
        }
        if(null == contentType || length > 0 && contentType.includes(MediaType.APPLICATION_FORM_URLENCODED)){
            return readFormData(exchange,chain,gatewayLog);
        }
        if (length <=0 || contentType.includes(MediaType.MULTIPART_FORM_DATA)){
            return readResponse(exchange,chain,gatewayLog);
        }
        return chain.filter(exchange);
    }

    private Mono<Void> readResponse(ServerWebExchange exchange, GatewayFilterChain chain, GatewayLog gatewayLog) {
        ServerHttpResponse responseDecorator = logResponse(exchange, gatewayLog,chain);
        return chain.filter(exchange.mutate().response(responseDecorator).build());
    }


    /**
     * ReadFormData  读取表单数据
     * @param exchange
     * @param chain
     * @return
     */
    private Mono<Void> readFormData(ServerWebExchange exchange, GatewayFilterChain chain, GatewayLog gatewayLog){
        HttpHeaders headers = exchange.getRequest().getHeaders();
        if (headers.containsKey("Sec-WebSocket-Key")){
            return chain.filter(exchange);
        }

        return exchange.getFormData()
                .doOnNext(formData -> {
                    gatewayLog.setFormData(formData);
                    log.info("[GatewayLogFilter](Request) FormData {}",formData);
                }).then(Mono.defer(() -> {
                    MultiValueMap<String, String> formData = gatewayLog.getFormData();
                    Charset charset = null;
                    try {
                        charset = headers.getContentType().getCharset();
                    } catch (Exception e) {
                    }
                    charset = charset == null? StandardCharsets.UTF_8:charset;
                    StringBuilder formDataBodyBuilder = new StringBuilder();
                    try {
                        /*
                         * repackage form data
                         */
                        for (Map.Entry<String, List<String>> entry : formData.entrySet()) {
                            for(String value : entry.getValue()){
                                formDataBodyBuilder.append(entry.getKey()).append("=").append(URLEncoder.encode(value, charset.name())).append("&");
                            }
                        }
                    }catch (UnsupportedEncodingException e){}
                    /*
                     * substring with the last char '&'
                     */
                    String formDataBodyString = "";
                    if(formDataBodyBuilder.length()>0){
                        formDataBodyString = formDataBodyBuilder.substring(0, formDataBodyBuilder.length() - 1);
                    }
                    /*
                     * get data bytes
                     */
                    byte[] bodyBytes =  formDataBodyString.getBytes(charset);
                    int contentLength = bodyBytes.length;
                    HttpHeaders httpHeaders = new HttpHeaders();
                    httpHeaders.putAll(exchange.getRequest().getHeaders());
                    httpHeaders.remove(HttpHeaders.CONTENT_LENGTH);
                    /*
                     * in case of content-length not matched
                     */
                    httpHeaders.setContentLength(contentLength);
                    /*
                     * use BodyInserter to InsertFormData Body
                     */
                    BodyInserter<String, ReactiveHttpOutputMessage> bodyInserter = BodyInserters.fromObject(formDataBodyString);
                    CachedBodyOutputMessage cachedBodyOutputMessage = new CachedBodyOutputMessage(exchange, httpHeaders);
                    if (log.isDebugEnabled()){
                        log.debug("[GatewayLogFilter]Rewrite Form Data :{}",formDataBodyString);
                    }
                    return bodyInserter.insert(cachedBodyOutputMessage,  new BodyInserterContext())
                            .then(Mono.defer(() -> {
                                ServerHttpRequestDecorator decorator = new ServerHttpRequestDecorator(exchange.getRequest()) {
                                    @Override
                                    public HttpHeaders getHeaders() {
                                        return httpHeaders;
                                    }
                                    @Override
                                    public Flux<DataBuffer> getBody() {
                                        return cachedBodyOutputMessage.getBody();
                                    }
                                };
                                ServerHttpResponse responseDecorator = logResponse(exchange, gatewayLog, chain);
                                return chain.filter(exchange.mutate().request(decorator).response(responseDecorator).build());
                            }));
                }));
    }

    private Mono<Void> handleRequestBody(ServerWebExchange exchange, GatewayFilterChain chain, GatewayLog gatewayLog){
        ServerRequest serverRequest =  ServerRequest.create(exchange, HandlerStrategies.withDefaults().messageReaders());

        Mono<String> modifyReqBody = serverRequest.bodyToMono(String.class).flatMap(reqBody -> {
            String modifyBody = handleRequestBody(exchange, reqBody);
            gatewayLog.setRequestBody(modifyBody);
            return Mono.just(modifyBody);
        });

        //重写请求体后需要重写content length
        HttpHeaders headers = new HttpHeaders();
        headers.putAll(exchange.getRequest().getHeaders());
        headers.remove(HttpHeaders.CONTENT_LENGTH);

        BodyInserter bodyInserter = BodyInserters.fromPublisher(modifyReqBody, String.class);
        CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange, headers);

        return bodyInserter.insert(outputMessage, new BodyInserterContext()).then(Mono.defer(() -> {
            ServerHttpRequest decorator = decorate(exchange, headers, outputMessage);
            ServerHttpResponse responseDecorator = logResponse(exchange, gatewayLog, chain);
            return chain.filter(exchange.mutate().request(decorator).response(responseDecorator).build());
        }));
    }

    private ServerHttpRequestDecorator decorate(ServerWebExchange exchange, HttpHeaders headers, CachedBodyOutputMessage outputMessage) {
        return new ServerHttpRequestDecorator(exchange.getRequest()) {
                    @Override
                    public HttpHeaders getHeaders() {
                        long contentLength = headers.getContentLength();
                        HttpHeaders httpHeaders = new HttpHeaders();
                        httpHeaders.putAll(super.getHeaders());
                        if (contentLength > 0) {
                            httpHeaders.setContentLength(contentLength);
                        } else {
                            httpHeaders.set(HttpHeaders.TRANSFER_ENCODING, "chunked");
                        }
                        return httpHeaders;
                    }

                    @Override
                    public Flux<DataBuffer> getBody() {
                        return outputMessage.getBody();
                    }
                };
    }

    private String handleRequestBody(ServerWebExchange exchange,String body) {
        if (requestBodyHandlers != null && !requestBodyHandlers.isEmpty()){
            for (RequestBodyHandler requestBodyHandler : requestBodyHandlers) {
                body = requestBodyHandler.handler(exchange,body);
            }
        }
        return body;
    }

    /**
     * log response exclude response body
     * @param exchange
     * @param gatewayLog
     * @param chain
     */
    private ServerHttpResponse logResponse(ServerWebExchange exchange, GatewayLog gatewayLog, GatewayFilterChain chain){
        ServerHttpResponse response = exchange.getResponse();
        if (!gatewayProperties.getLogResponse()){
            return response;
        }
        //重写原始响应
        return  new BodyHandlerServerHttpResponseDecorator(initBodyHandler(exchange, gatewayLog,chain), response);
    }
    /**
     * 响应body处理，添加响应的打印
     *
     * @param exchange
     * @param gatewayLog
     * @param chain
     * @return
     */
    protected BodyHandlerFunction initBodyHandler(ServerWebExchange exchange, GatewayLog gatewayLog, GatewayFilterChain chain) {
        return (resp, body) -> {
            HttpHeaders orgHeaders = resp.getHeaders();
            HttpHeaders headers = HttpHeadersFilter.filter(headersFilters,orgHeaders, exchange, HttpHeadersFilter.Type.RESPONSE);
            gatewayLog.setResponseHeaders(headers).setResponseStatus(resp.getStatusCode());
            //拦截
            MediaType respMedia = headers.getContentType();
            if (respMedia != null && respMedia.includes(MediaType.APPLICATION_JSON)){
                //拦截
                String originalResponseContentType = exchange.getAttribute(ORIGINAL_RESPONSE_CONTENT_TYPE_ATTR);
                HttpHeaders httpHeaders = new HttpHeaders();
                httpHeaders.setContentType(MediaType.parseMediaType(originalResponseContentType));

                DefaultClientResponseAdapter clientResponseAdapter = new DefaultClientResponseAdapter(body, httpHeaders);

                Class inClass = Object.class;
                Class outClass = Object.class;
                Mono mono = clientResponseAdapter.bodyToMono(inClass);
                Mono modifiedBody = mono.flatMap(originalBody -> {
                    try {
                        gatewayLog.setResponseBody(objectMapper.writeValueAsString(originalBody));
                    } catch (Exception e) {}
                    return rewriteResponseHandler.apply(exchange, originalBody);
                });

                BodyInserter bodyInserter = BodyInserters.fromPublisher(modifiedBody, outClass);
                CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange, exchange.getResponse().getHeaders());
                return bodyInserter.insert(outputMessage, new BodyInserterContext())
                        .then(Mono.defer(() -> {
                            Flux<DataBuffer> messageBody = outputMessage.getBody();
                            if (!headers.containsKey(HttpHeaders.TRANSFER_ENCODING)) {
                                messageBody = messageBody.doOnNext(data -> headers.setContentLength(data.readableByteCount()));
                            }
                            return resp.writeWith(messageBody);
                        }));
            }else {
               return chain.filter(exchange);
            }
        };
    }
}